# Copyright (c) 2024-present AI-Labs

from fastapi import APIRouter

import base64
import os
import torch
import uuid

from datetime import datetime

from funasr import AutoModel
from funasr.utils.postprocess_utils import rich_transcription_postprocess

from configs import config

"""
路由信息设置
"""
router = APIRouter(
    prefix='/asr/funasr',
    tags = ['语音识别']
)

"""
模型信息加载
"""
model = AutoModel(
    model="iic/SenseVoiceSmall",
    vad_model="fsmn-vad",
    vad_kwargs={"max_single_segment_time": 30000},
    device=config.service.funasr.device,
)


"""
自动语音识别，将语音信息转换成文本数据
"""
@router.post("/v1")
def audio_to_text(request: dict):
    # 先将语音数据保存成本地音频文件
    localdir = f"upload/{datetime.now().strftime('%Y-%m-%d')}"
    os.makedirs(f"{config.setting.statics.path}/{localdir}", exist_ok=True)
    localfile = f"{config.setting.statics.path}/{localdir}/{uuid.uuid4()}"
    with open(localfile, "wb") as f:
        f.write(base64.b64decode(request["audio"]))

    # 模型对本地音频文件进行识别，生成文本内容
    res = model.generate(
        input=localfile,
        cache={},
        language="auto",  # "zn", "en", "yue", "ja", "ko", "nospeech"
        use_itn=True,
        batch_size_s=60,
        merge_vad=True,  #
        merge_length_s=15,
    )

    result = rich_transcription_postprocess(res[0]["text"])

    # 清理显存缓存
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.ipc_collect()

    # 返回识别结果
    return {"asr_result": result}
